from p4utils.utils.helper import load_topo
from p4utils.utils.sswitch_thrift_API import SimpleSwitchThriftAPI

class Controller:
    
    def __init__(self):
        self.topo = load_topo('topology.json')
        self.controllers = {}
        self.init()
    
    def init(self):
        self.connect_to_switches()
        self.reset_states()
        self.set_table_defaults()
    
    def connect_to_switches(self):
        for p4switch in self.topo.get_p4switches():
            thrift_port = self.topo.get_thrift_port(p4switch)
            self.controllers[p4switch] = SimpleSwitchThriftAPI(thrift_port)
    
    def reset_states(self):
        for controller in self.controllers.values():
            controller.reset_state()
    
    def set_table_defaults(self):
        for controller in self.controllers.values():
            # 设置默认动作
            controller.table_set_default("ipv4_lpm", "drop", [])
            controller.table_set_default("src_route_forward", "drop", [])
            controller.table_set_default("first_hop_classifier", "NoAction", [])
    
    def add_ipv4_route(self, switch_name, dst_ip, prefix_len, next_hop_mac, egress_port):
        self.controllers[switch_name].table_add(
            "ipv4_lpm", "ipv4_forward",
            [f"{dst_ip}/{prefix_len}"],
            [next_hop_mac, str(egress_port)]
        )
    
    def add_src_route_path(self, switch_name, src_mac, dst_ip, path_len, ports):
        self.controllers[switch_name].table_add(
            "first_hop_classifier", "add_src_route_path",
            [src_mac, dst_ip],
            [str(path_len)] + [str(p) for p in ports[:4]]
        )
    
    def add_src_route_forward(self, switch_name, in_port, out_port):
        self.controllers[switch_name].table_add(
            "src_route_forward", "src_route_forward",
            [str(in_port)],
            [str(out_port)]
        )
    
    def program_switch_rules(self):
        # 示例配置 - 需要根据实际拓扑修改
        # 配置IPv4路由
        self.add_ipv4_route("s1", "10.0.1.1", 32, "00:00:00:00:01:01", 1)
        self.add_ipv4_route("s1", "10.0.2.2", 32, "00:00:00:00:02:02", 2)
        
        # 配置源路由路径
        self.add_src_route_path("s1", "00:00:00:00:01:01", "10.0.3.3", 3, [2, 3, 1])
        
        # 配置源路由转发规则
        self.add_src_route_forward("s1", 1, 2)
        self.add_src_route_forward("s1", 2, 3)
        self.add_src_route_forward("s1", 3, 1)

if __name__ == "__main__":
    controller = Controller()
    controller.program_switch_rules()